package cn.com.pc.content.mongo;

import com.mongodb.client.result.DeleteResult;
import org.springframework.dao.OptimisticLockingFailureException;
import org.springframework.data.domain.*;
import org.springframework.data.mongodb.core.MongoOperations;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.repository.query.MongoEntityInformation;
import org.springframework.data.mongodb.repository.support.SimpleMongoRepository;
import org.springframework.data.repository.support.PageableExecutionUtils;
import org.springframework.data.util.StreamUtils;
import org.springframework.data.util.Streamable;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import static org.springframework.data.mongodb.core.query.Criteria.where;

public class CustomMongoRepositoryImpl<T, ID> implements CustomMongoRepository<T, ID>{

    private final MongoOperations mongoOperations;
    private final MongoEntityInformation<T, ID> entityInformation;
    private String collectionName = "";

    @Override
    public String getCollectionName() {
        return this.collectionName;
    }

    @Override
    public void setCollectionName(String collectionName) {
        this.collectionName = collectionName;
    }

    /**
     * Creates a new {@link SimpleMongoRepository} for the given {@link MongoEntityInformation} and {@link MongoTemplate}.
     *
     * @param metadata must not be {@literal null}.
     * @param mongoOperations must not be {@literal null}.
     */
    public CustomMongoRepositoryImpl(MongoEntityInformation<T, ID> metadata, MongoOperations mongoOperations) {

        Assert.notNull(metadata, "MongoEntityInformation must not be null!");
        Assert.notNull(mongoOperations, "MongoOperations must not be null!");

        this.entityInformation = metadata;
        this.mongoOperations = mongoOperations;
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#save(java.lang.Object)
     */
    @Override
    public <S extends T> S save(S entity) {

        Assert.notNull(entity, "Entity must not be null!");

        if (entityInformation.isNew(entity)) {
            return mongoOperations.insert(entity, entityInformation.getCollectionName());
        }

        return mongoOperations.save(entity, entityInformation.getCollectionName());
    }

    @Override
    public <S extends T> S save(S entity, String collectionName) {

        Assert.notNull(entity, "Entity must not be null!");

        if (entityInformation.isNew(entity)) {
            return mongoOperations.insert(entity, collectionName);
        }

        return mongoOperations.save(entity, collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#saveAll(java.lang.Iterable)
     */
    @Override
    public <S extends T> List<S> saveAll(Iterable<S> entities) {

        Assert.notNull(entities, "The given Iterable of entities not be null!");

        Streamable<S> source = Streamable.of(entities);
        boolean allNew = source.stream().allMatch(entityInformation::isNew);

        if (allNew) {

            List<S> result = source.stream().collect(Collectors.toList());
            return new ArrayList<>(mongoOperations.insert(result, entityInformation.getCollectionName()));
        }

        return source.stream().map(this::save).collect(Collectors.toList());
    }

    @Override
    public <S extends T> List<S> saveAll(Iterable<S> entities, String collectionName) {

        Assert.notNull(entities, "The given Iterable of entities not be null!");

        Streamable<S> source = Streamable.of(entities);
        boolean allNew = source.stream().allMatch(entityInformation::isNew);

        if (allNew) {

            List<S> result = source.stream().collect(Collectors.toList());
            return new ArrayList<>(mongoOperations.insert(result, collectionName));
        }

        return source.stream().map(this::save).collect(Collectors.toList());
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#findById(java.io.Serializable)
     */
    @Override
    public Optional<T> findById(ID id) {

        Assert.notNull(id, "The given id must not be null!");

        return Optional.ofNullable(
                mongoOperations.findById(id, entityInformation.getJavaType(), entityInformation.getCollectionName()));
    }

    @Override
    public Optional<T> findById(ID id, String collectionName) {

        Assert.notNull(id, "The given id must not be null!");

        return Optional.ofNullable(
                mongoOperations.findById(id, entityInformation.getJavaType(), collectionName));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#existsById(java.lang.Object)
     */
    @Override
    public boolean existsById(ID id) {

        Assert.notNull(id, "The given id must not be null!");

        return mongoOperations.exists(getIdQuery(id), entityInformation.getJavaType(),
                entityInformation.getCollectionName());
    }

    @Override
    public boolean existsById(ID id, String collectionName) {

        Assert.notNull(id, "The given id must not be null!");

        return mongoOperations.exists(getIdQuery(id), entityInformation.getJavaType(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#count()
     */
    @Override
    public long count() {
        return mongoOperations.count(new Query(), entityInformation.getCollectionName());
    }

    @Override
    public long count(String collectionName) {
        return mongoOperations.count(new Query(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#deleteById(java.lang.Object)
     */
    @Override
    public void deleteById(ID id) {

        Assert.notNull(id, "The given id must not be null!");

        mongoOperations.remove(getIdQuery(id), entityInformation.getJavaType(), entityInformation.getCollectionName());
    }

    @Override
    public void deleteById(ID id, String collectionName) {

        Assert.notNull(id, "The given id must not be null!");

        mongoOperations.remove(getIdQuery(id), entityInformation.getJavaType(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#delete(java.lang.Object)
     */
    @Override
    public void delete(T entity) {

        Assert.notNull(entity, "The given entity must not be null!");

        DeleteResult deleteResult = mongoOperations.remove(entity, entityInformation.getCollectionName());

        if (entityInformation.isVersioned() && deleteResult.wasAcknowledged() && deleteResult.getDeletedCount() == 0) {
            throw new OptimisticLockingFailureException(String.format(
                    "The entity with id %s with version %s in %s cannot be deleted! Was it modified or deleted in the meantime?",
                    entityInformation.getId(entity), entityInformation.getVersion(entity),
                    entityInformation.getCollectionName()));
        }
    }

    @Override
    public void delete(T entity, String collectionName) {

        Assert.notNull(entity, "The given entity must not be null!");

        DeleteResult deleteResult = mongoOperations.remove(entity, collectionName);

        if (entityInformation.isVersioned() && deleteResult.wasAcknowledged() && deleteResult.getDeletedCount() == 0) {
            throw new OptimisticLockingFailureException(String.format(
                    "The entity with id %s with version %s in %s cannot be deleted! Was it modified or deleted in the meantime?",
                    entityInformation.getId(entity), entityInformation.getVersion(entity),
                    collectionName));
        }
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#delete(java.lang.Iterable)
     */
    @Override
    public void deleteAll(Iterable<? extends T> entities) {

        Assert.notNull(entities, "The given Iterable of entities not be null!");

        entities.forEach(this::delete);
    }

    @Override
    public void deleteAll(Iterable<? extends T> entities,  String collectionName) {

        Assert.notNull(entities, "The given Iterable of entities not be null!");

        entities.forEach(entity -> this.delete(entity, collectionName));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#deleteAll()
     */
    @Override
    public void deleteAll() {
        mongoOperations.remove(new Query(), entityInformation.getCollectionName());
    }

    @Override
    public void deleteAll(String collectionName) {
        mongoOperations.remove(new Query(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#findAll()
     */
    @Override
    public List<T> findAll() {
        return findAll(new Query());
    }

    @Override
    public List<T> findAll(String collectionName) {
        return findAll(new Query(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.CrudRepository#findAllById(java.lang.Iterable)
     */
    @Override
    public Iterable<T> findAllById(Iterable<ID> ids) {

        Assert.notNull(ids, "The given Ids of entities not be null!");

        return findAll(new Query(new Criteria(entityInformation.getIdAttribute())
                .in(Streamable.of(ids).stream().collect(StreamUtils.toUnmodifiableList()))));
    }

    @Override
    public Iterable<T> findAllById(Iterable<ID> ids, String collectionName) {

        Assert.notNull(ids, "The given Ids of entities not be null!");

        return findAll(new Query(new Criteria(entityInformation.getIdAttribute())
                .in(Streamable.of(ids).stream().collect(StreamUtils.toUnmodifiableList()))), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.PagingAndSortingRepository#findAll(org.springframework.data.domain.Pageable)
     */

    @Override
    public Page<T> findAll(Pageable pageable) {

        Assert.notNull(pageable, "Pageable must not be null!");

        long count = count();
        List<T> list = findAll(new Query().with(pageable));

        return new PageImpl<>(list, pageable, count);
    }

    @Override
    public Page<T> findAll(Pageable pageable, String collectionName) {

        Assert.notNull(pageable, "Pageable must not be null!");

        long count = count();
        List<T> list = findAll(new Query().with(pageable), collectionName);

        return new PageImpl<>(list, pageable, count);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.PagingAndSortingRepository#findAll(org.springframework.data.domain.Sort)
     */
    @Override
    public List<T> findAll(Sort sort) {

        Assert.notNull(sort, "Sort must not be null!");

        return findAll(new Query().with(sort));
    }

    @Override
    public List<T> findAll(Sort sort, String collectionName) {

        Assert.notNull(sort, "Sort must not be null!");

        return findAll(new Query().with(sort), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#insert(java.lang.Object)
     */
    @Override
    public <S extends T> S insert(S entity) {

        Assert.notNull(entity, "Entity must not be null!");

        return mongoOperations.insert(entity, entityInformation.getCollectionName());
    }

    @Override
    public <S extends T> S insert(S entity, String collectionName) {

        Assert.notNull(entity, "Entity must not be null!");

        return mongoOperations.insert(entity, collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#insert(java.lang.Iterable)
     */
    @Override
    public <S extends T> List<S> insert(Iterable<S> entities) {

        Assert.notNull(entities, "The given Iterable of entities not be null!");

        List<S> list = Streamable.of(entities).stream().collect(StreamUtils.toUnmodifiableList());

        if (list.isEmpty()) {
            return list;
        }

        return new ArrayList<>(mongoOperations.insertAll(list));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#findAllByExample(org.springframework.data.domain.Example, org.springframework.data.domain.Pageable)
     */
    @Override
    public <S extends T> Page<S> findAll(final Example<S> example, Pageable pageable) {

        Assert.notNull(example, "Sample must not be null!");
        Assert.notNull(pageable, "Pageable must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation()).with(pageable); //

        List<S> list = mongoOperations.find(query, example.getProbeType(), entityInformation.getCollectionName());

        return PageableExecutionUtils.getPage(list, pageable,
                () -> mongoOperations.count(Query.of(query).limit(-1).skip(-1), example.getProbeType(), entityInformation.getCollectionName()));
    }

    @Override
    public <S extends T> Page<S> findAll(final Example<S> example, Pageable pageable, String collectionName) {

        Assert.notNull(example, "Sample must not be null!");
        Assert.notNull(pageable, "Pageable must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation()).with(pageable); //

        List<S> list = mongoOperations.find(query, example.getProbeType(), collectionName);

        return PageableExecutionUtils.getPage(list, pageable,
                () -> mongoOperations.count(Query.of(query).limit(-1).skip(-1), example.getProbeType(), collectionName));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#findAllByExample(org.springframework.data.domain.Example, org.springframework.data.domain.Sort)
     */
    @Override
    public <S extends T> List<S> findAll(Example<S> example, Sort sort) {

        Assert.notNull(example, "Sample must not be null!");
        Assert.notNull(sort, "Sort must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation()) //
                .with(sort);

        return mongoOperations.find(query, example.getProbeType(), entityInformation.getCollectionName());
    }

    @Override
    public <S extends T> List<S> findAll(Example<S> example, Sort sort, String collectionName) {

        Assert.notNull(example, "Sample must not be null!");
        Assert.notNull(sort, "Sort must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation()) //
                .with(sort);

        return mongoOperations.find(query, example.getProbeType(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.repository.MongoRepository#findAllByExample(org.springframework.data.domain.Example)
     */
    @Override
    public <S extends T> List<S> findAll(Example<S> example) {
        return findAll(example, Sort.unsorted());
    }

    @Override
    public <S extends T> List<S> findAll(Example<S> example, String collectionName) {
        return findAll(example, Sort.unsorted(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.query.QueryByExampleExecutor#findOne(org.springframework.data.domain.Example)
     */
    @Override
    public <S extends T> Optional<S> findOne(Example<S> example) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return Optional
                .ofNullable(mongoOperations.findOne(query, example.getProbeType(), entityInformation.getCollectionName()));
    }

    @Override
    public <S extends T> Optional<S> findOne(Example<S> example, String collectionName) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return Optional
                .ofNullable(mongoOperations.findOne(query, example.getProbeType(), collectionName));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.query.QueryByExampleExecutor#count(org.springframework.data.domain.Example)
     */
    @Override
    public <S extends T> long count(Example<S> example) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return mongoOperations.count(query, example.getProbeType(), entityInformation.getCollectionName());
    }

    @Override
    public <S extends T> long count(Example<S> example, String collectionName) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return mongoOperations.count(query, example.getProbeType(), collectionName);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.query.QueryByExampleExecutor#exists(org.springframework.data.domain.Example)
     */
    @Override
    public <S extends T> boolean exists(Example<S> example) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return mongoOperations.exists(query, example.getProbeType(), entityInformation.getCollectionName());
    }

    @Override
    public <S extends T> boolean exists(Example<S> example, String collectionName) {

        Assert.notNull(example, "Sample must not be null!");

        Query query = new Query(new Criteria().alike(example)) //
                .collation(entityInformation.getCollation());

        return mongoOperations.exists(query, example.getProbeType(), collectionName);
    }

    private Query getIdQuery(Object id) {
        return new Query(getIdCriteria(id));
    }

    private Criteria getIdCriteria(Object id) {
        return where(entityInformation.getIdAttribute()).is(id);
    }

    private List<T> findAll(@Nullable Query query) {

        if (query == null) {
            return Collections.emptyList();
        }

        return mongoOperations.find(query, entityInformation.getJavaType(), entityInformation.getCollectionName());
    }

    private List<T> findAll(@Nullable Query query, String collectionName) {

        if (query == null) {
            return Collections.emptyList();
        }

        return mongoOperations.find(query, entityInformation.getJavaType(), collectionName);
    }
}

